package com.xs.langchain4j_springboot.splitter;

import dev.langchain4j.internal.ValidationUtils;

import java.util.function.Function;

public class CustomizedSegmentBuilder {

    private final int maxSegmentSize;
    private final Function<String, Integer> sizeFunction;
    private final String joinSeparator;
    private final int joinSeparatorSize;
    private String segment = "";
    private int segmentSize = 0;

    public CustomizedSegmentBuilder(int maxSegmentSize, Function<String, Integer> sizeFunction, String joinSeparator) {
        this.maxSegmentSize = ValidationUtils.ensureGreaterThanZero(maxSegmentSize, "maxSegmentSize");
        this.sizeFunction = (Function)ValidationUtils.ensureNotNull(sizeFunction, "sizeFunction");
        this.joinSeparator = (String)ValidationUtils.ensureNotNull(joinSeparator, "joinSeparator");
        this.joinSeparatorSize = this.sizeOf(joinSeparator);
    }

    public int getSize() {
        return this.segmentSize;
    }

    public boolean hasSpaceFor(String text) {
        int totalSize = this.sizeOf(text);
        if (this.isNotEmpty()) {
            totalSize += this.segmentSize + this.joinSeparatorSize;
        }

        return totalSize <= this.maxSegmentSize;
    }

    public boolean hasSpaceFor(int size) {
        int totalSize = size;
        if (this.isNotEmpty()) {
            totalSize = size + this.segmentSize + this.joinSeparatorSize;
        }

        return totalSize <= this.maxSegmentSize;
    }

    public int sizeOf(String text) {
        return (Integer)this.sizeFunction.apply(text);
    }

    public void append(String text) {
        if (this.isNotEmpty()) {
            this.segment = this.segment + this.joinSeparator;
        }

        this.segment = this.segment + text;
        this.segmentSize = this.sizeOf(this.segment);
    }

    public void prepend(String text) {
        if (this.isNotEmpty()) {
            this.segment = text + this.joinSeparator + this.segment;
        } else {
            this.segment = text;
        }

        this.segmentSize = this.sizeOf(this.segment);
    }

    public boolean isNotEmpty() {
        return !this.segment.isEmpty();
    }

    public String toString() {
        return this.segment.trim();
    }

    public void reset() {
        this.segment = "";
        this.segmentSize = 0;
    }
}
